// 引入 TensorFlow.js
import * as tf from '@tensorflow/tfjs';
import * as tfvis from '@tensorflow/tfjs-vis';
// 准备数据: y = x^2 + 2x + 3
function generateData(numSamples) {
    const xs = [];
    const ys = [];
    for (let i = 0; i < numSamples; i++) {
        const x = i;
        const y = x * x + 2 * x + 3;
        xs.push([x]);
        ys.push([y]);
    }
    return { xs, ys };
}

const { xs, ys } = generateData(10); // 生成100个样本

// 转换为张量
const inputTensor = tf.tensor2d(xs);
const labelTensor = tf.tensor2d(ys);
console.log(xs, ys);
tfvis.render.scatterplot(
    { name: 'Xs vs Ys' },
    { values: xs.map((x, i) => ({ x, y: ys[i] })) },
    {
        xAxisDomain: [-1, 10],
        yAxisDomain: [2, 103]
    }
);
// 构建模型
const model = tf.sequential();
model.add(tf.layers.dense({ inputShape: [1], units: 300, activation: 'relu' }));
model.add(tf.layers.dense({ units: 1 }));

model.compile({
    optimizer: tf.train.adam(),
    loss: tf.losses.meanSquaredError,
    metrics: ['mse'],
});

// 训练模型
async function trainModel() {
    const epochs = 500;
    const history = await model.fit(inputTensor, labelTensor, {
        epochs: epochs,
        batchSize: 100,// 批次大小
        validationSplit: 0.1,// 验证集
        callbacks: tfvis.show.fitCallbacks(
            { name: '训练过程' },
            ['loss', 'val_loss'],
            { callbacks: ['onEpochEnd'] }
        )
    });
    console.log('训练完成');
    return history;
}

trainModel().then(history => {
    console.log('训练损失:', history.history.loss);
    console.log('验证损失:', history.history.val_loss);

    // 使用模型进行预测
    model.predict(tf.tensor([[0]])).print();
});